import math
import argparse
import curses
from copy import deepcopy
from dgo3 import GoBoard  # 从 dgo3.py 中导入 GoBoard 类
import time
from QLAgent import QLearningAgent, train_agent

# 继承 GoBoard 类，并添加 reset 方法
class CustomGoBoard(GoBoard):

    def __init__(self, stdscr, size):
        # 调用父类 GoBoard 的 __init__ 方法
        # 由于 GoBoard 的 __init__ 方法需要 file_name 参数，你可以传递一个占位符或实际文件名
        super().__init__(stdscr, file_name="game_output5.txt")  # 使用 "dummy.sgf" 作为占位符文件名
        self.size = size  # 添加 size 属性

    def make_move(self, move, player):
        if self.board[move[0]][move[1]] == 0:
            self.board[move[0]][move[1]] = player
            return True
        return False

    def reset(self):
        self.board = [[0] * 19 for _ in range(19)]
        self.current_player = 1
        self.print_board()

    def check_draw(self):
        for row in range(self.size):
            for col in range(self.size):
                if self.board[row][col] == 0:
                    return False
        return True

    def check_three_in_a_row(self, player):
        # 检查行
        for row in range(self.size):
            for col in range(self.size - 2):
                if all(self.board[row][col + i] == player for i in range(3)):
                    return True
        # 检查列
        for col in range(self.size):
            for row in range(self.size - 2):
                if all(self.board[row + i][col] == player for i in range(3)):
                    return True
        # 检查对角线（左上到右下）
        for row in range(self.size - 2):
            for col in range(self.size - 2):
                if all(self.board[row + i][col + i] == player for i in range(3)):
                    return True
        # 检查对角线（右上到左下）
        for row in range(self.size - 2):
            for col in range(2, self.size):
                if all(self.board[row + i][col - i] == player for i in range(3)):
                    return True
        return False

    def check_four_in_a_row(self, player):
        # 检查行
        for row in range(self.size):
            for col in range(self.size - 3):
                if all(self.board[row][col + i] == player for i in range(4)):
                    return True
        # 检查列
        for col in range(self.size):
            for row in range(self.size - 3):
                if all(self.board[row + i][col] == player for i in range(4)):
                    return True
        # 检查对角线（左上到右下）
        for row in range(self.size - 3):
            for col in range(self.size - 3):
                if all(self.board[row + i][col + i] == player for i in range(4)):
                    return True
        # 检查对角线（右上到左下）
        for row in range(self.size - 3):
            for col in range(3, self.size):
                if all(self.board[row + i][col - i] == player for i in range(4)):
                    return True
        return False

# 定义环境类
class Environment:
    def __init__(self, size=19):
        self.size = size
        self.stdscr = curses.initscr()
        curses.start_color()
        curses.init_pair(1, curses.COLOR_BLACK, curses.COLOR_WHITE)
        curses.init_pair(2, curses.COLOR_WHITE, curses.COLOR_BLACK)
        self.board = CustomGoBoard(self.stdscr, size=size)  # 使用继承的 CustomGoBoard 类
        self.state = [[0] * size for _ in range(size)]
        self.current_player = 1
        self.episode_rewards = []

    def reset(self):
        self.board.reset()  # 调用 CustomGoBoard 的 reset 方法
        self.state = [[0] * self.size for _ in range(self.size)]
        self.current_player = 1
        return self.state

    def step(self, action):
        row, col = action
        if self.board.make_move(action, self.current_player):
            self.board.print_board(False)
            # time.sleep(0.1)

            reward = 0.1  # 默认合法动作奖励

            if self.board.check_win(self.current_player):
                reward = 10  # 赢棋奖励
                done = True
            elif self.board.check_draw():
                reward = 0.5  # 平局奖励
                done = True
            else:
                # 检查是否形成三连、四连等
                if self.board.check_three_in_a_row(self.current_player):
                    reward += 0.5
                if self.board.check_four_in_a_row(self.current_player):
                    reward += 1
                # 检查是否被对手形成三连、四连
                opponent = -self.current_player
                if self.board.check_three_in_a_row(opponent):
                    reward -= 0.5
                if self.board.check_four_in_a_row(opponent):
                    reward -= 1
                done = False
                self.current_player = -self.current_player
        else:
            reward = -1  # 非法动作惩罚
            done = False
        self.state[row][col] = self.current_player
        return self.state, reward, done

    def render(self):
        self.board.print_board(False)

    def close(self):
        curses.endwin()

    def check_space(self):
        """当用户按下空格键时，调出棋盘"""
        self.board.print_board(True)
        key = self.stdscr.getch()
        while key != ord(' '):  # 等待用户按下空格键
            key = self.stdscr.getch()
        self.board.print_board(False)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Gobang Game")
    parser.add_argument("--size", type=int, default=19, choices=range(5, 20), help="Size of the game board (5-19)")
    args = parser.parse_args()
    board_size = args.size

    agent = QLearningAgent(size=board_size)
    env = Environment(size=board_size)
    train_agent(agent, env, episodes=1000)